import os

import cv2
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision.transforms import functional as F
import random

class MyDataset(Dataset):
    def __init__(self, dataset_dir: str, dataset_type: str, device='cuda'):
        super(MyDataset, self).__init__()
        self.device = device
        self.dataset_dir = dataset_dir
        self.dataset_type = dataset_type
        with open(os.path.join(self.dataset_dir, "{}.txt".format(self.dataset_type))) as f:
            self.data_names = f.read().splitlines()

    def __getitem__(self, index):
        img_path = os.path.join(self.dataset_dir, 'images', self.data_names[index] + '.jpg')
        mask_path = os.path.join(self.dataset_dir, 'masks', self.data_names[index] + '.png')
        img_np = cv2.imread(img_path)
        img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
        mask_tensor = F.to_tensor(Image.open(mask_path).convert('1'))
        data = {
            'img': img_np,
            'mask': mask_tensor
        }
        # 显卡或者cpu
        if self.device == 'cuda':
            for key in data.keys():
                if isinstance(data[key], torch.Tensor):
                    data[key] = data[key].cuda()

        data['data_name'] = self.data_names[index]
        return data

    def __len__(self):
        return len(self.data_names)

    def shuffle(self):
        random.shuffle(self.data_names)



if __name__ == "__main__":
    dataset = MyDataset(r'Datasets/Crack500', 'train', 'cuda')
    a = dataset[0]
    print(a)
    print(len(dataset))
